# -*- coding: utf-8 -*-
"""
@Time : 2024/5/31 2:16 PM
@Author: imamhusan
@Des: 描述
"""
import torch
from transformers import BertTokenizer, BertForSequenceClassification
from torch.utils.data import DataLoader, Dataset
from torch.optim import AdamW

# سانلىق مەلۇماتلار (مىسال ئۈچۈن)
train_texts = [
    "بۇ بىر ياخشى كىتاب",
    "مەن بۇ فىلىمنى ياقتۇرمىدىم",
    "ھاۋا بۈگۈن بەك گۈزەل",
    "بۇ تاماق تەمسىز"
]
train_labels = [1, 0, 1, 0]  # 1: ئىجابىي، 0: سەلبىي

# BERT مودېلى ۋە tokenizer نى يۈكلەش
model_name = "bert-base-multilingual-cased"  # كۆپ تىللىق BERT مودېلى
tokenizer = BertTokenizer.from_pretrained(model_name)
model = BertForSequenceClassification.from_pretrained(model_name, num_labels=2)

# سانلىق مەلۇمات توپلىمى كلاسسى
class TextDataset(Dataset):
    def __init__(self, texts, labels):
        self.texts = texts
        self.labels = labels

    def __len__(self):
        return len(self.texts)

    def __getitem__(self, idx):
        return self.texts[idx], self.labels[idx]

# سانلىق مەلۇماتلارنى ئالدىن بىر تەرەپ قىلىش
def collate_fn(batch):
    texts, labels = zip(*batch)
    encoded = tokenizer(texts, padding=True, truncation=True, return_tensors="pt")
    return encoded, torch.tensor(labels)

# سانلىق مەلۇمات يۈكلىگۈچنى قۇرۇش
dataset = TextDataset(train_texts, train_labels)
dataloader = DataLoader(dataset, batch_size=2, shuffle=True, collate_fn=collate_fn)

# مەشىقلەندۈرۈش پارامېتىرلىرى
optimizer = AdamW(model.parameters(), lr=5e-5)
num_epochs = 3

# مەشىقلەندۈرۈش ھالقىسى
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)

for epoch in range(num_epochs):
    model.train()
    for batch in dataloader:
        inputs, labels = batch
        inputs = {k: v.to(device) for k, v in inputs.items()}
        labels = labels.to(device)

        optimizer.zero_grad()
        outputs = model(**inputs, labels=labels)
        loss = outputs.loss
        loss.backward()
        optimizer.step()

    print(f"Epoch {epoch+1}/{num_epochs}, Loss: {loss.item():.4f}")

# يېڭى تېكىستنى سىناش
def predict(text):
    model.eval()
    with torch.no_grad():
        inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True)
        inputs = {k: v.to(device) for k, v in inputs.items()}
        outputs = model(**inputs)
        probabilities = torch.softmax(outputs.logits, dim=1)
        prediction = torch.argmax(probabilities, dim=1).item()
        return "ئىجابىي" if prediction == 1 else "سەلبىي"

# سىناش
test_sentence = "بۇ كىتاب ناھايىتى ياخشى"
print(f"'{test_sentence}' دېگەن جۈملىنىڭ پوزىتسىيەسى: {predict(test_sentence)}")